/***************************************************************************\
 * The Mattise Kernel														*
 * Copyright 2007 Matthew Iselin											*
 * Licensed under the GPL													*
 *																			*
 * ata_pio.cc																*
 *																			*
 * Port I/O ATA block driver.												*
 * Currently handles drive detection and the low-level features as well.	*
 *																			*
\***************************************************************************/

#include <system.h>
#include <memory.h>
#include <console.h>
#include <time.h>
#include <iostream>
#include <disk.h>
#include <stdio.h>
#include <io.h>
#include <string.h>

// ATA ports (offset from base)
// write
#define ATA_FEATS	1
#define ATA_SCOUNT	2
#define ATA_SNUM	3
#define ATA_CLOW	4
#define ATA_CHI		5
#define ATA_HDEV	6
#define ATA_COMM	7
// read
#define ATA_DATA	0
#define ATA_ERR		1
#define ATA_STAT	7

// drive types
#define DRIVE_ATA	0
#define DRIVE_ATAPI	1

// drive to host short
#define dtohs(n) ( (((n) & 0xFF00) >> 8) | (((n) & 0x00FF) << 8) )

// list of drives
Drive drvlist[32];

// waiters
volatile int prictl_Wait = 1;
volatile int secctl_Wait = 1;

// gets a drive
Drive GetDrive( int drv )
{
	return drvlist[drv];
}

// waits for a command to finish, after sending it
void waitCommand( int ctl, Drive* drv, int offset, unsigned char command )
{
	// primary?
	if( ctl == 0 )
	{
		// we will be waiting
		prictl_Wait = 1;
		
		// send it
		outportb( drv->base + offset, command );
		
		// read in the status byte
		uchar_t ident_stat = inportb( drv->base + ATA_COMM );
		if( ident_stat & 0x1 )
			return;
		
		// wait!
		while( prictl_Wait );
	}
	
	// secondary controller
	if( ctl == 1 )
	{
		// we will be waiting
		secctl_Wait = 1;
		
		// send it
		outportb( drv->base + offset, command );
		
		// read in the status byte
		uchar_t ident_stat = inportb( drv->base + ATA_COMM );
		if( ident_stat & 0x1 )
			return;
		
		// wait!
		while( secctl_Wait );
	}
}

// IRQ14: primary hard drive controller
void hdd_IRQ14( struct regs* r )
{
	// command finished
	prictl_Wait = 0;
}

// IRQ15: secondary hard drive controller
void hdd_IRQ15( struct regs* r )
{
	// command finished
	secctl_Wait = 0;
}

// detects type of drive - after an 'identify device'
int GetDriveType( Drive* drv )
{
	// read in the registers 'count' and 'sector'
	unsigned char sc = inportb( drv->base + ATA_SCOUNT ), sn = inportb( drv->base + ATA_SNUM );
	
	cons.Output( "[IDE ] Detecting ATA device type... ", true );
	
	// check...
	if( sc == 1 && sn == 1 )
	{
		// ATAPI
		cons.Output( "ATAPI\n", true );
		return DRIVE_ATAPI;
	}
	
	// ATA
	cons.Output( "ATA\n", true );
	return DRIVE_ATA;
}

// executes the 'Identify Device' command
void IdentifyDevice( Drive* drv )
{	
	// choose the drive
	outportb( drv->base + ATA_HDEV, 0xE0 | drv->drv );
	
	// send & wait
	waitCommand( 0, drv, ATA_COMM, 0xEC );
	//outportb( drv->base + ATA_COMM, 0xEC );
	
	// wait until we're ready
	while( (inportb( drv->base + 0x206 ) & 0x80) );
	
	// read in the status byte
	uchar_t ident_stat = inportb( drv->base + ATA_COMM );
	
	// error?
	if( ident_stat & 0x1 )
	{
		// TODO: handle packet device identification, page 105 of the ATA specs
		kprintf( "IDENTIFY_DEVICE FAILED!\n" );
		return;
	}
	
	// detect type
	drv->type = GetDriveType( drv );
	
	// check
	if( drv->type == DRIVE_ATAPI )
	{
		return;
	}
	
	// wait for the data to arrive - but timeout if it doesn't
	if( ( inportb( drv->base + ATA_COMM ) & 0x08 ) )
	{
		// counter
		int idx;
		
		// read in all the data
		unsigned short* identdata = (unsigned short*) &drv->ident;
		for( idx = 0; idx < 255; idx++ )
		{
			identdata[idx] = inportw( drv->base );
		}
	}
}

// finds a drive
bool FindDrive( short port, char drv )
{
	// temporary data goes here
	uchar_t tmp;
	
	// check...
	outportb( port, drv );
	
	// wait a bit
	//delay( 4 );
	
	// read the status port
	tmp = inportb( port + 1 );
	
	// busy?
	return( tmp & 0x40 );
}

// prologue to any disk operation - CHS
void chs_prolog( int addr, Drive* drv )
{
	// convert to CHS first
	uchar_t Sector = ( addr % drv->ident.sectors ) + 1;
	ushort_t Cylinder = ( addr / drv->ident.sectors ) / drv->ident.heads;
	uchar_t Head = ( addr / drv->ident.sectors ) % drv->ident.heads;
	
	// null to features
	outportb( drv->base + ATA_FEATS, 0x00 );
	
	// number of sectors
	outportb( drv->base + ATA_SCOUNT, 0x01 );
	
	// block address
	outportb( drv->base + ATA_SNUM, Sector );
	outportb( drv->base + ATA_CLOW, (uchar_t) (Cylinder & 0xFF) );
	outportb( drv->base + ATA_CHI, (uchar_t) ((Cylinder & 0xFF00)>>16) );
	
	// drive id et al
	outportb( drv->base + ATA_HDEV, ( drv->drv << 4 ) | ( Head & 0x0F ) );
}

// prologue to any disk operation - all except the command itself (LBA28)
void lba28_prolog( int addr, Drive* drv )
{
	// wait until we're ready
	while( (inportb( drv->base + 0x206 ) & 0x80) );
	
	// null to features
	outportb( drv->base + ATA_FEATS, 0x00 );
	
	// number of sectors
	outportb( drv->base + ATA_SCOUNT, 0x01 );
	
	// block address
	outportb( drv->base + ATA_SNUM, (uchar_t) addr );
	outportb( drv->base + ATA_CLOW, (uchar_t) (addr >> 8) );
	outportb( drv->base + ATA_CHI, (uchar_t) (addr >> 16) );

	// drive id et al
	uchar_t hdev = 0xE0 | ( drv->drv << 4 ) | ( addr >> 24 );
	outportb( drv->base + ATA_HDEV, hdev );
}

// prologue to any disk operation - all except the command itself (LBA48)
void lba48_prolog( unsigned long long addr, Drive* drv )
{
	// nulls to features
	outportb( drv->base + ATA_FEATS, 0x00 );
	outportb( drv->base + ATA_FEATS, 0x00 );
	
	// number of sectors
	outportb( drv->base + ATA_SCOUNT, 0x00 );
	outportb( drv->base + ATA_SCOUNT, 0x01 );
	
	// block address
	outportb( drv->base + ATA_SNUM, (uchar_t) (addr >> 24) );
	outportb( drv->base + ATA_SNUM, (uchar_t) addr );
	outportb( drv->base + ATA_CLOW, (uchar_t) (addr >> 32) );
	outportb( drv->base + ATA_CLOW, (uchar_t) (addr >> 8) );
	outportb( drv->base + ATA_CHI, (uchar_t) (addr >> 40) );
	outportb( drv->base + ATA_CHI, (uchar_t) (addr >> 16) );
	
	// drive id et al
	outportb( drv->base + ATA_HDEV, 0x40 | ( drv->drv << 4 ) );
}

// starts a read (CHS)
void chs_startread( int lba_addr, Drive* drv )
{
	// start the op
	chs_prolog( lba_addr, drv );

	// read command
	waitCommand( 0, drv, ATA_COMM, 0x20 );
}

// starts a write (CHS)
void chs_startwrite( int lba_addr, Drive* drv )
{
	// start the op
	chs_prolog( lba_addr, drv );

	// read command
	waitCommand( 0, drv, ATA_COMM, 0x30 );
}

// starts a read (LBA28)
void lba28_startread( int lba_addr, Drive* drv )
{	
	// start the op
	lba28_prolog( lba_addr, drv );

	// read command
	waitCommand( 0, drv, ATA_COMM, 0x20 );
//	outportb( drv->base + ATA_COMM, 0x20 );
	
	// wait until we're ready
	while( (inportb( drv->base + ATA_COMM ) & 0x80) );
}

// starts a write (LBA28)
void lba28_startwrite( int lba_addr, Drive* drv )
{
	// start the op
	lba28_prolog( lba_addr, drv );

	// read command
	waitCommand( 0, drv, ATA_COMM, 0x30 );
}

// starts a read (LBA48)
void lba48_startread( unsigned long long lba_addr, Drive* drv )
{
	// start the op
	lba48_prolog( lba_addr, drv );

	// read command
	waitCommand( 0, drv, ATA_COMM, 0x24 );
}

// starts a write (LBA48)
void lba48_startwrite( unsigned long long lba_addr, Drive* drv )
{
	// start the op
	lba48_prolog( lba_addr, drv );

	// read command
	waitCommand( 0, drv, ATA_COMM, 0x34 );
}

// read a sector of data (LBA28)
void lba28_ReadSector( int drv, int block, char* dat )
{
	// start the read
	lba28_startread( block, &drvlist[drv] );
	
	// check for validity
	uchar_t errreg = inportb( drvlist[drv].base + 0x206 ) & 0x1;
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait for bytes
	while( ! ( inportb( drvlist[drv].base + 0x206 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	
	// 'null count'
	int nc = 0;
	
	// for debugging purposes
	#define ATA_LBA28_READ_TYPE	ushort_t
	#define ATA_LBA28_INPORT( d ) inportw( d )
	
	// read in the data
	memset( dat, 0, 512 );
	for( idx = 0; idx < drvlist[drv].ident.sector_bytes; idx += sizeof( ATA_LBA28_READ_TYPE ) )
	{
		ATA_LBA28_READ_TYPE tmpword = ATA_LBA28_INPORT( drvlist[drv].base );
		*((ATA_LBA28_READ_TYPE*) &dat[idx]) = tmpword;
	}
}

// write a sector of data (LBA28)
void lba28_WriteSector( int drv, int block, char* dat )
{	
	// start the read
	lba28_startread( block, &drvlist[drv] );
	
	// check for validity
	uchar_t errreg = inportb( drvlist[drv].base + ATA_ERR );
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait for bytes
	while( ! ( inportb( drvlist[drv].base + 7 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	for( idx = 0; idx < 512; idx += sizeof( short ) )
	{
		outportw( drvlist[drv].base, *((short*)&dat[idx]) );
	}
}

// read a sector of data (LBA48)
void lba48_ReadSector( unsigned long long drv, int block, char* dat )
{	
	// start the read
	lba48_startread( block, &drvlist[drv] );
	
	// check for validity
	uchar_t errreg = inportb( drvlist[drv].base + ATA_ERR );
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait for bytes
	while( ! ( inportb( drvlist[drv].base + 7 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	for( idx = 0; idx < 512; idx += sizeof( short ) )
	{
		short tmpword = inportw( drvlist[drv].base );
		*((short*) &dat[idx]) = tmpword;
	}
}

// write a sector of data (LBA48)
void lba48_WriteSector( int drv, unsigned long long block, char* dat )
{	
	// start the read
	lba48_startwrite( block, &drvlist[drv] );
	
	// check for validity
	uchar_t errreg = inportb( drvlist[drv].base + ATA_ERR );
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait for bytes
	while( ! ( inportb( drvlist[drv].base + 7 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	for( idx = 0; idx < 512; idx += sizeof( short ) )
	{
		outportw( drvlist[drv].base, *((short*)&dat[idx]) );
	}
}

// read a sector of data (CHS)
void chs_ReadSector( int drv, unsigned long long block, char* dat )
{	
	// start the read
	chs_startread( block, &drvlist[drv] );
	
	// check for validity
	uchar_t errreg = inportb( drvlist[drv].base + ATA_ERR );
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait for bytes
	while( ! ( inportb( drvlist[drv].base + 7 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	for( idx = 0; idx < 512; idx += sizeof( short ) )
	{
		short tmpword = inportw( drvlist[drv].base );
		*((short*) &dat[idx]) = tmpword;
	}
}

// write a sector of data (CHS)
void chs_WriteSector( int drv, int block, char* dat )
{	
	// start the read
	chs_startwrite( block, &drvlist[drv] );
	
	// check for validity
	uchar_t errreg = inportb( drvlist[drv].base + ATA_ERR );
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait for bytes
	while( ! ( inportb( drvlist[drv].base + 7 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	for( idx = 0; idx < 512; idx += sizeof( short ) )
	{
		outportw( drvlist[drv].base, *((short*)&dat[idx]) );
	}
}

// hard drive read/write (it's almost the same anyway)
int hdd_rw( int fd, void* buff, uint_t count, bool isread )
{
	// valid descriptor?
	if( fd == -1 )
		return 0;
	
	// get the device number and drive number
	int i = GetDeviceID( fd );
	fd = GetDeviceNumber( fd );
	
	// char pointer to the buffer
	char* ptr = (char*) buff;
	
	// read in all the sectors desired
	for( int z = 0; z < count; z++ )
	{
		// read in the sector
		if( drvlist[i].lba )
		{
			// > 120 GB = use LBA48
			if( drvlist[i].ident.lba_capacity > 251658240 ) /** 251658240 sectors in 120 GiB **/
			{
				if( isread )
					lba48_ReadSector( drvlist[i].id, drvlist[i].offset + z, ptr );
				else
					lba48_WriteSector( drvlist[i].id, drvlist[i].offset + z, ptr );
			}
			else
			{
				if( isread )
					lba28_ReadSector( drvlist[i].id, drvlist[i].offset + z, ptr );
				else
					lba28_WriteSector( drvlist[i].id, drvlist[i].offset + z, ptr );
			}
		}
		else
		{
			// CHS
			if( isread )
				chs_ReadSector( drvlist[i].id, drvlist[i].offset + z, ptr );
			else
				chs_WriteSector( drvlist[i].id, drvlist[i].offset + z, ptr );
		}
		
		// the drive uses 512-byte sectors
		ptr += 512;
	}
	
	// increment the offset
	drvlist[i].offset += count;
	
	// all done
	return count;
}

// hard drive READ - BLOCK read...
int hdd_read( int fd, void* buff, uint_t count )
{
	// read it in
	return hdd_rw( fd, buff, count, true );
}

// hard drive WRITE - BLOCK write...
int hdd_write( int fd, void* buff, uint_t count )
{
	// read it in
	return hdd_rw( fd, buff, count, false );
}

// hard drive lseek
int hdd_lseek( int fd, uint_t offset, int where )
{
	// where are we going from?
	if( where == SEEK_SET )
	{
		drvlist[GetDeviceID(fd)].offset = offset;
	}
	else if( where == SEEK_END )
	{
		drvlist[GetDeviceID(fd)].offset = drvlist[GetDeviceID(fd)].ident.lba_capacity - offset;
	}
	else if( where == SEEK_CUR )
	{
		drvlist[GetDeviceID(fd)].offset += offset;
	}
	else
	{
		return 0;
	}
	return offset;
}

// installs a new drive
void InstallDrive( int i )
{
	// the device name
	char* devname = (char*) kmalloc( 5 );
	sprintf( devname, "dr%d", i );
	
	// install the new device - the device ID is the drive number
	int d = InstallDevice( const_cast< char* >( devname ), 0, i );
	
	// free the memory
	kfree( devname );
	
	// install it
	drvlist[i].devicenum = d;
	drvlist[i].lba = ( drvlist[i].ident.capability & 0x02 );
	
	// install the functions
	INSTFUNCDATA ataDev = { (IOCtlIntFunc) NULL, (IOCtlPtrFunc) NULL, hdd_read, hdd_write, hdd_lseek, (GetCFunc) NULL, (PutCFunc) NULL };
	InstallFunctions( d, ataDev );
}

// detect drives
void DetectDrives()
{
	// first find out which controllers are active
	int IDEPri = 0;
	int IDESec = 0;
	int myIndex = 1;
	uchar_t byt = 0;

	// start with primary controller
	outportb( 0x1F3, 0x88 );
	byt = inportb( 0x1F3 );
	
	// exists?
	if( byt == 0x88 )
		IDEPri = 1;
	else
		IDEPri = 0;
		
	// then secondary controller
	outportb( 0x173, 0x88 );
	byt = inportb( 0x173 );
	
	// exists?
	if( byt == 0x88 )
		IDESec = 1;
	else
		IDESec = 0;

	// index
	int i = 0;
	
	// actual drives
	if( IDEPri )
	{
		// install the irq handler (IRQ14)
		irq_install_handler( 14, hdd_IRQ14 );
		
		// find the primary master and slave
		if( FindDrive( 0x1F6, 0xA0 ) )
		{
			cons.Output( "[IDE ] Primary master found\n", true );
			drvlist[i].id = i;
			drvlist[i].base = 0x1F0;
			drvlist[i].type = 0;
			drvlist[i].drv = /*0x80 +*/ i;
			IdentifyDevice( &drvlist[i] );
			InstallDrive( i );
			i++;
		}
		if( FindDrive( 0x1F6, 0xB0 ) )
		{
			cons.Output( "[IDE ] Primary slave found\n", true );
			drvlist[i].id = i;
			drvlist[i].base = 0x1F0;
			drvlist[i].type = 0;
			drvlist[i].drv = /*0x80 +*/ i;
			IdentifyDevice( &drvlist[i] );
			InstallDrive( i );
			i++;
		}
	}
	if( IDESec )
	{
		// install the irq handler (IRQ15)
		irq_install_handler( 15, hdd_IRQ15 );
		
		// find the secondary master and slave
		if( FindDrive( 0x176, 0xA0 ) )
		{
			cons.Output( "[IDE ] Secondary master found\n", true );
			drvlist[i].id = i;
			drvlist[i].base = 0x170;
			drvlist[i].type = 0;
			drvlist[i].drv = /*0x80 +*/ i;
			IdentifyDevice( &drvlist[i] );
			InstallDrive( i );
			i++;
		}
		if( FindDrive( 0x176, 0xB0 ) )
		{
			cons.Output( "[IDE ] Secondary slave found\n", true );
			drvlist[i].id = i;
			drvlist[i].base = 0x170;
			drvlist[i].type = 0;
			drvlist[i].drv = /*0x80 +*/ i;
			IdentifyDevice( &drvlist[i] );
			InstallDrive( i );
			i++;
		}
	}
}
